import numpy as np
import pandas as pd
import talib
import traceback


def tdxFunc(num,data1,data2):
    return func_map.get(num, lambda x, y: [0] * len(x))(data1, data2)

def tdxFuncA(num,dl):
    return func_map.get(num, lambda x: [0] * x)(dl)


#--------------------------------------------------------------------------------
# 每个函数最好都使用try,不然一旦函数出错股票程序就会崩溃。

df = pd.DataFrame()

#把O和C两组数据保存到全局变量df中
def save_OC(listO, listC):
    global df
    try:
        df = pd.DataFrame()
        print(listO[:5])
        df['O'] = listO
        print(df['O'].iloc[:10])
        df['C'] = listC
    except Exception as e:
        msg = traceback.format_exc()
        print(msg)

def save_HL(listH, listL):
    global df
    try:
        df['H'] = listH
        df['L'] = listL
    except Exception as e:
        msg = traceback.format_exc()
        print(msg)

def save_VA(listV, listA):
    global df
    try:
        df['V'] = listV
        df['A'] = listA
    except Exception as e:
        msg = traceback.format_exc()
        print(msg)



def calculate_sma(series, n, m):
    """
    计算 SMA (Simple Moving Average)
    参数:
    series: pandas Series, 输入数据
    n: int, 移动平均的周期
    m: int, 权重
    返回:
    pandas Series, SMA 结果
    """
    sma = np.zeros_like(series)
    sma[:n] = np.nan
    
    for i in range(n, len(series)):
        if i == n:
            sma[i] = series[i-n+1:i+1].mean()
        else:
            sma[i] = (series[i] * m + sma[i-1] * (n-m)) / n
    
    return pd.Series(sma, index=series.index)


# 计算K的值
def calc_KDJ(listM1, listM2):
    global df
    try:
        # 获取参数,由于KDJ不需要Open开盘数据，这里把df['O'] 用来保参数N了
        N = int(df['O'].iloc[0])
        M1 = int(listM1[0])
        M2 = int(listM2[0])

        MinL = df['L'].rolling(window=N, min_periods=1).min()
        MaxH = df['H'].rolling(window=N, min_periods=1).max()
        df['RSV'] = (df['C']-MinL)/(MaxH-MinL)*100

        df['K'] = calculate_sma(df['RSV'], n=M1, m=1)
        df['D'] = calculate_sma(df['K'], n=M2, m=1)
        df['J'] = 3*df['K'] - 2*df['D']

        return df['K'].tolist()

    except Exception as e:
        msg = traceback.format_exc()
        print(msg)
        return np.zeros_like(lista).tolist()


def return_D(dl):
    global df
    try:
        return df['D'].tolist()
    except Exception as e:
        msg = traceback.format_exc()
        print(msg)
        return np.zeros((dl,),dtype=float).tolist()


def return_J(dl):
    try:
        return df['J'].tolist()
    except Exception as e:
        msg = traceback.format_exc()
        print(msg)
        return np.zeros((dl,),dtype=float).tolist()






#在func_map这里注册函数为几号函数，供通达信公式调用。
#save_开头的函数，都对应DLL 2号函数有两个列表参数 不用返回值
#calc_开头的函数，都对应DLL 1号函数有两个列表参数 需要返回值
#return_开头的函数，都对应DLL 3号函数有1个int参数代表返回值的长度 需要返回值

# func_map要放在程序最后面
func_map = {
    #1，2，3几个函数专门用来保存数据，
    1: save_OC, #保存数据Open和Close
    2: save_HL, #保存数据High和Low
    3: save_VA, #保存数据Vol和Amount

    4: calc_KDJ,  #计算KDJ中的K值 
    5: return_D,  #计算KDJ中的D值 
    6: return_J,  #计算KDJ中的J值 
}    